from copy import deepcopy
from network import QKDNetwork
from docplex.mp.model import Model

def get_m_by_i(i: int, n: int):
    i_tmp = deepcopy(i)
    n_tmp = deepcopy(n)
    start = -1
    while(i_tmp >= 0):
        n_tmp -= 1
        i_tmp -= n_tmp
        start += 1
    return start

def lp_with_pre(net: QKDNetwork):
    n = len(net.nodeList)
    model = Model(name='LP_with_pre')
    
    i_count = net.sd_num
    j_count = 5
    k_count = 5
    n_count = net.maxCandidatePathLength - 1
    l_count = net.numLinks
    m_count = len(net.nodeList)
    
    p = model.continuous_var_dict(
            [(i, j, n, k)  for i in range(i_count)
                        for j in range(j_count)
                        for n in range(n_count)
                        for k in range(k_count)], lb=0, ub=1, name='p'
        )
    
    x = model.binary_var_dict(
            [(i, j)  for i in range(i_count)
                        for j in range(j_count)], name='x'
        )
    
    _alpha = model.continuous_var(lb=0, name='alpha')
    
    _lambda = model.continuous_var_dict(
            [(m, l)  for m in range(m_count)
                        for l in range(l_count)], lb=0, name='lambda'
        )
    
    # R_mk lambda_ml <= Ra + a/t
    for l in range(l_count):
        flow_volume = model.linear_expr()
        flow_volume -= net.getLinkCapacityByLinkID(l) * _alpha
        for m in range(m_count):
            for k in net.getNeighborNodeIDs(m):
                flow_volume += net.getLinkCapacity(m, k) * _lambda[m, l]
        model.add_constraint(flow_volume <= 0)
    
    # sum p_ij^nl + x <= lambda_ml
    for i in range(i_count):
        m = get_m_by_i(i)
        for l in range(l_count):
            flow_volume_pass = model.linear_expr()
            flow_volume_pass -= _lambda[m, l]
            for j in range(j_count):
                flow_volume_pass += x[i, j] * net.isIJLeverageL(i, j, l)
                for n in range(n_count):
                    for k in range(k_count):
                        flow_volume_pass += p[i, j, n, k] * net.isIJNLeverageL(i, j, n, k, l)
            model.add_constraint(flow_volume_pass <= 0)
            
    # sum x = 1
    for i in range(i_count):
        sole_path = model.linear_expr()
        for j in range(j_count):
            sole_path += x[i, j]
        model.add_constraint(sole_path == 1)
        
    # sum p = 1
    for i in range(i_count):
        for j in range(j_count):
            for n in range(n_count):
                fix_flow = model.linear_expr()
                for k in range(k_count):
                    fix_flow += p[i, j, n, k]
                model.add_constraint(sole_path == 1)
                
    model.minimize(_alpha)
            
    sol = model.solve()
    if sol:
        print("alpha:", sol.objective_value)
    